{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from itertools import *\n",
    "from functools import *\n",
    "import operator\n",
    "from scipy.stats import norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# 最大值函数\n",
    "def I(x: np.ndarray):\n",
    "    t = np.zeros([x.shape[0], 2])\n",
    "    t[:, 0] = x\n",
    "    return np.max(t, axis=1)\n",
    "\n",
    "# 期望函数\n",
    "def E(x: np.ndarray, p: np.ndarray):\n",
    "    return np.dot(x, p)\n",
    "\n",
    "# 依照正态分布计算概率值\n",
    "def p_norm(x: np.ndarray, mu, sigma):\n",
    "    init_arr = norm.cdf(x, mu, sigma)\n",
    "    res_arr = np.zeros(x.shape[0])\n",
    "    res_arr[0] = init_arr[0]\n",
    "    res_arr[1:] = np.diff(init_arr)\n",
    "    res_arr[-1] = 1 - init_arr[-1]\n",
    "    return res_arr\n",
    "\n",
    "# Variable\n",
    "T = 3\n",
    "alpha = 0.9\n",
    "# Dm\n",
    "mu_m = 2\n",
    "sigma_m = 0.1\n",
    "D_m = np.array(list(range(0, 6)))\n",
    "p_D_m = p_norm(D_m, mu_m, sigma_m)\n",
    "# Dt\n",
    "mu_t = 5\n",
    "sigma_t = 0.5\n",
    "D_t = np.array(list(range(0, 10)))\n",
    "p_D_t = p_norm(D_t, mu_t, sigma_t)\n",
    "p = 10\n",
    "c_e = 20\n",
    "c_m = 2\n",
    "h_e = 5\n",
    "h_m = 1\n",
    "theta = h_m / (1 - alpha)\n",
    "b = 2\n",
    "\n",
    "initial = True"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "# 以下为 t = T 时的代码, 可以看出, 与任意 t 时刻的不同之处在于 psi 处的 f 函数取值为 0(初值条件)\n",
    "def psi_1T(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z):\n",
    "    return h_e * E(I(v_1_bar - e - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - v_1_bar + e), p_D_t) \\\n",
    "           + c_e * (v_1_bar - e) \\\n",
    "           + c_m * (z - v_4_bar) \\\n",
    "           + h_m * (v_4_bar - v_1_bar) \\\n",
    "           + theta * e\n",
    "\n",
    "def psi_2T(y, v_2_bar, v_3_bar, v_4_bar, z):\n",
    "    return h_e * E(I(y - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - y), p_D_t) \\\n",
    "           + c_e * y \\\n",
    "           + c_m * (z - v_4_bar) \\\n",
    "           + h_m * (v_4_bar - y)\n",
    "\n",
    "def psi_3T(y, v_3_bar, v_4_bar, z):\n",
    "    return h_e * E(I(y - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - y), p_D_t) \\\n",
    "           + c_e * y \\\n",
    "           + c_m * (z - v_4_bar) \\\n",
    "           + h_m * (v_4_bar - y)\n",
    "\n",
    "def psi_4T(y, v_4_bar, z):\n",
    "    return h_e * E(I(y - D_t), p_D_t) \\\n",
    "           + b * E(I(D_t - y), p_D_t) \\\n",
    "           + c_e * y \\\n",
    "           + c_m * (z - v_4_bar) \\\n",
    "           + h_m * (v_4_bar - y)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "def phi_1T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    e_range = range(0, v_1_bar - v_0_bar +1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    e_z_grid = product(e_range, z_range)\n",
    "    psi_1T_values = [psi_1T(v_1_bar, v_2_bar, v_3_bar,v_4_bar, e, z) for e, z in e_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_1T_values)\n",
    "\n",
    "def phi_2T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_1_bar, v_2_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_2T_values = [psi_2T(y, v_2_bar, v_3_bar, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_2T_values)\n",
    "\n",
    "def phi_3T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_2_bar, v_3_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_3T_values = [psi_3T(y, v_3_bar, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_3T_values)\n",
    "\n",
    "def phi_4T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_3_bar, v_4_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_4T_values = [psi_4T(y, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_4T_values)\n",
    "\n",
    "def g_T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    return np.min([\n",
    "        phi_1T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_2T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_3T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_4T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar)\n",
    "    ])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "def f_T(v_0, v_1, v_2, v_3, v_4):\n",
    "    _D_m = []\n",
    "    for D_m_i in D_m:\n",
    "        exp_1_value = np.min([p * (D_m_i - a) + g_T(v_0, v_1 - a, v_2 - a, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_1 - v_0]) + 1)])\n",
    "        exp_2_value = np.min([p * (D_m_i - a) + g_T(v_0, v_0, v_2 - a, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_2 - v_0]) + 1)])\n",
    "        exp_3_value = np.min([p * (D_m_i - a) + g_T(v_0, v_0, v_0, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_3 - v_0]) + 1)])\n",
    "        exp_4_value = np.min([p * (D_m_i - a) + g_T(v_0, v_0, v_0, v_0, v_4 - a) for a in range(0, np.min([D_m_i, v_4 - v_0]) + 1)])\n",
    "        _D_m.append(np.min([exp_1_value, exp_2_value, exp_3_value, exp_4_value]))\n",
    "    return E(_D_m, p_D_m)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "# 以下为任意 t 时刻的代码\n",
    "def psi_1t(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z):\n",
    "    global initial\n",
    "    val = psi_1T(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z)\n",
    "    if initial:\n",
    "       return val\n",
    "\n",
    "    _D_t = []\n",
    "    for D_t_i in D_t:\n",
    "        _D_t.append(f_T(v_1_bar - e - D_t_i, v_2_bar - e - D_t_i, v_3_bar - e - D_t_i, v_4_bar - e -D_t_i, z - e -D_t_i))\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "\n",
    "def psi_2t(y, v_2_bar, v_3_bar, v_4_bar, z):\n",
    "    global initial\n",
    "    val = psi_2T(y, v_2_bar, v_3_bar, v_4_bar, z)\n",
    "    if initial:\n",
    "       return val\n",
    "\n",
    "    _D_t = []\n",
    "    for D_t_i in D_t:\n",
    "        _D_t.append(f_T(y - D_t_i, v_2_bar - D_t_i, v_3_bar - D_t_i, v_4_bar - D_t_i, z - D_t_i))\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "\n",
    "def psi_3t(y, v_3_bar, v_4_bar, z):\n",
    "    global initial\n",
    "    val = psi_3T(y, v_3_bar, v_4_bar, z)\n",
    "    if initial:\n",
    "       return val\n",
    "\n",
    "    _D_t = []\n",
    "    for D_t_i in D_t:\n",
    "        _D_t.append(f_T(y - D_t_i, y - D_t_i, v_3_bar - D_t_i, v_4_bar - D_t_i, z - D_t_i))\n",
    "    return val + alpha * E(_D_t, p_D_t)\n",
    "\n",
    "def psi_4t(y, v_4_bar, z):\n",
    "    global initial\n",
    "    val = psi_4T(y, v_4_bar, z)\n",
    "    if initial:\n",
    "       return val\n",
    "\n",
    "    _D_t = []\n",
    "    for D_t_i in D_t:\n",
    "        _D_t.append(f_T(y - D_t_i, y - D_t_i, y -D_t_i, v_4_bar - D_t_i, z - D_t_i))\n",
    "    return val + alpha * E(_D_t, p_D_t)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "def phi_1t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    e_range = range(0, v_1_bar - v_0_bar +1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    e_z_grid = product(e_range, z_range)\n",
    "    psi_1t_values = [psi_1t(v_1_bar, v_2_bar, v_3_bar,v_4_bar, e, z) for e, z in e_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_1t_values)\n",
    "\n",
    "def phi_2t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_1_bar, v_2_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_2t_values = [psi_2t(y, v_2_bar, v_3_bar, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_2t_values)\n",
    "\n",
    "def phi_3t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_2_bar, v_3_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_3t_values = [psi_3t(y, v_3_bar, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_3t_values)\n",
    "\n",
    "def phi_4t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    y_range = range(v_3_bar, v_4_bar + 1)\n",
    "    z_range = range(v_4_bar, v_4_bar + 21)\n",
    "    y_z_grid = product(y_range, z_range)\n",
    "    psi_4t_values = [psi_4t(y, v_4_bar, z) for y, z in y_z_grid]\n",
    "    return - c_e * v_0_bar + np.min(psi_4t_values)\n",
    "\n",
    "def g_t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):\n",
    "    return np.min([\n",
    "        phi_1t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_2t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_3t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),\n",
    "        phi_4t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar)\n",
    "    ])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "def f_t(v_0, v_1, v_2, v_3, v_4):\n",
    "    _D_m = []\n",
    "    for D_m_i in D_m:\n",
    "        exp_1_value = np.min([p * (D_m_i - a) + g_t(v_0, v_1 - a, v_2 - a, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_1 - v_0]) + 1)])\n",
    "        exp_2_value = np.min([p * (D_m_i - a) + g_t(v_0, v_0, v_2 - a, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_2 - v_0]) + 1)])\n",
    "        exp_3_value = np.min([p * (D_m_i - a) + g_t(v_0, v_0, v_0, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_3 - v_0]) + 1)])\n",
    "        exp_4_value = np.min([p * (D_m_i - a) + g_t(v_0, v_0, v_0, v_0, v_4 - a) for a in range(0, np.min([D_m_i, v_4 - v_0]) + 1)])\n",
    "        _D_m.append(np.min([exp_1_value, exp_2_value, exp_3_value, exp_4_value]))\n",
    "    return E(_D_m, p_D_m)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "data": {
      "text/plain": "10.499999984214604"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "initial = False\n",
    "f_t(1, 2, 3, 4, 5)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}